from PIL import Image, ImageDraw
import os

def mm_to_pixels(mm: float, resolution_m_per_pixel: float) -> int:
    """Convert millimeters to pixels given resolution in meters/pixel."""
    # 1 pixel represents resolution_m_per_pixel meters
    # 1 meter = 1000 millimeters
    return int(round((mm / 1000.0) / resolution_m_per_pixel))


def draw_grid_map(width_mm, height_mm, resolution_m_per_pixel, border_width_pixels, output_filename,
                  rect_obstacles_mm=None, square_obstacles_mm=None):
    """
    绘制二值化的栅格地图。

    参数:
    width_mm (int): 场地的物理宽度，单位为毫米。
    height_mm (int): 场地的物理高度，单位为毫米。
    resolution_m_per_pixel (float): 分辨率，单位为米/像素。
    border_width_pixels (int): 边界障碍物的宽度，单位为像素。
    output_filename (str): 输出的PGM文件路径。
    """
    
    # 毫米转换为米
    width_m = width_mm / 1000.0
    height_m = height_mm / 1000.0
    
    # 自由区域像素尺寸（不包含边框）
    free_width_pixels = int(round(width_m / resolution_m_per_pixel))
    free_height_pixels = int(round(height_m / resolution_m_per_pixel))

    # 总像素尺寸（包含边框）
    # 根据需求：删除右侧黑边，右侧边框视为 0 像素，仅在左/上/下保留边框
    width_pixels = free_width_pixels + border_width_pixels  # 左侧有边框，右侧无边框
    height_pixels = free_height_pixels + 2 * border_width_pixels  # 上下各一条边框
    
    print(f"场地物理尺寸: {width_m}m x {height_m}m")
    print(f"地图分辨率: {resolution_m_per_pixel} m/pixel")
    print(
        f"图像像素尺寸: {width_pixels} x {height_pixels} "
        f"(自由区: {free_width_pixels} x {free_height_pixels}, 左/上/下边框: {border_width_pixels}px, 右边框: 0px)"
    )

    # 创建一个新的灰度图像 (L模式，8位)
    img = Image.new('L', (width_pixels, height_pixels), color=0)
    draw = ImageDraw.Draw(img)

    # 先整体填充为黑色，然后在中间绘制自由区域白色矩形
    # 自由区域左上角坐标（包含左/上边框的偏移）
    free_left = border_width_pixels
    free_top = border_width_pixels
    # 自由区宽高严格使用 60x180（3m x 9m 对应的像素），右侧无黑边
    free_right = free_left + free_width_pixels - 1  # 由于右侧无边框，正好到图像最右侧
    free_bottom = free_top + free_height_pixels - 1
    draw.rectangle([(free_left, free_top), (free_right, free_bottom)], fill=255)

    # 绘制矩形障碍（单位：毫米）
    if rect_obstacles_mm:
        for idx, rect in enumerate(rect_obstacles_mm):
            # 物体坐标以自由区域左上角为原点，绘制时需要加边框偏移
            x_px = border_width_pixels + mm_to_pixels(rect['x_mm'], resolution_m_per_pixel)
            y_px = border_width_pixels + mm_to_pixels(rect['y_mm'], resolution_m_per_pixel)
            w_px = mm_to_pixels(rect['w_mm'], resolution_m_per_pixel)
            h_px = mm_to_pixels(rect['h_mm'], resolution_m_per_pixel)
            # 以左上角(x_mm, y_mm)为基准
            draw.rectangle([(x_px, y_px), (x_px + w_px - 1, y_px + h_px - 1)], fill=0)

    # 绘制正方形障碍（用于小球，单位：毫米）
    if square_obstacles_mm:
        for sq in square_obstacles_mm:
            x_center_px = border_width_pixels + mm_to_pixels(sq['cx_mm'], resolution_m_per_pixel)
            y_center_px = border_width_pixels + mm_to_pixels(sq['cy_mm'], resolution_m_per_pixel)
            side_px = mm_to_pixels(sq['side_mm'], resolution_m_per_pixel)
            half = side_px // 2
            draw.rectangle([(x_center_px - half, y_center_px - half),
                            (x_center_px + side_px - half - 1, y_center_px + side_px - half - 1)], fill=0)

    # 保存为PGM格式
    img.save(output_filename)
    print(f"地图已成功保存为 {output_filename}")


if __name__ == "__main__":
    # 场地参数（单位：毫米）
    arena_width_mm = 9000
    arena_height_mm = 3000

    # ROS Navigation Stack 参数
    map_resolution = 0.05  # 米/像素（一个像素代表0.05m）

    # 边界障碍物宽度（像素）
    border_pixels = 5

    # 依据示意图的障碍物（单位均为毫米）。
    # 说明：下述坐标原点在左上角，x 向右、y 向下。
    # 假设与取值：
    # - 棕色箱子均为 500x500mm 正方形
    # - 右侧 3 个箱子从 x=5000mm 起，每个箱子之间水平间距 500mm（箱子之间留 500mm 空隙）
    # - 左侧 2 个箱子靠近左边界，x=200mm，纵向间隔约 1000mm
    # - 彩色球直径按 100mm 处理；上方球距上边界 400mm，对称设置下方球距下边界 400mm
    # 如需严格复现竞赛场地，请据实修改下面的数值。

    rectangle_obstacles_mm = [
        # 左侧两箱：左边缘与 y 轴重合 (x_left=0)，中心 y 分别为 1000mm / 2000mm
        {"x_mm": 0,   "y_mm": 750,  "w_mm": 500, "h_mm": 500},  # center_y=1000mm
        {"x_mm": 0,   "y_mm": 1750, "w_mm": 500, "h_mm": 500}, # center_y=2000mm
        # 右侧三箱（位于场地右半区）
        # 按照要求：最右侧箱子顶边贴上边界(y=0)，另外两只箱子的顶边距上边界分别为 750mm、1500mm
        {"x_mm": 5000, "y_mm": 1500, "w_mm": 500, "h_mm": 500},
        {"x_mm": 6000, "y_mm": 750,  "w_mm": 500, "h_mm": 500},
        {"x_mm": 7000, "y_mm": 0,    "w_mm": 500, "h_mm": 500},
    ]

    # 小球改为正方形表示
    ball_side_mm = 100
    top_margin_mm = 0
    bottom_margin_mm = arena_height_mm - top_margin_mm
    square_obstacles_mm = [
        # 顶部两处
        {"cx_mm": 2000, "cy_mm": top_margin_mm,    "side_mm": ball_side_mm},
        {"cx_mm": 4000, "cy_mm": top_margin_mm,    "side_mm": ball_side_mm},
        # 底部两处
        {"cx_mm": 2000, "cy_mm": bottom_margin_mm, "side_mm": ball_side_mm},
        {"cx_mm": 4000, "cy_mm": bottom_margin_mm, "side_mm": ball_side_mm},
    ]

    # 输出文件到脚本同目录，避免工作目录影响
    script_dir = os.path.dirname(os.path.abspath(__file__))
    output_map_filename = os.path.join(script_dir, "arena_map.pgm")

    draw_grid_map(
        arena_width_mm,
        arena_height_mm,
        map_resolution,
        border_pixels,
        output_map_filename,
        rect_obstacles_mm=rectangle_obstacles_mm,
        square_obstacles_mm=square_obstacles_mm,
    )